import pandas as pd
import numpy as np
import scipy.interpolate as spi
import matplotlib.pyplot as plt
import matplotlib
font = {'family': 'SimHei'}
matplotlib.rc('font', **font)

class Insert:
    def __init__(self):
        self.df = pd.read_table('test.txt', sep='   ', encoding='gbk', engine='python')

        self.df['#No. date'] = pd.DataFrame([''] + list(map(change_type, self.df['#No. date'])))
        plt.scatter(range(1, 31), self.df[' all'], label='全国确诊')
        plt.scatter(range(1, 31), self.df[' hubei'], label='湖北确诊')
        plt.scatter(range(1, 31), self.df['other'], label='湖北外确诊')
        plt.xticks(rotation=60)
        plt.xlabel('日期')
        plt.ylabel('人数')

        self.X = np.linspace(1, 31, 30)
        self.dic = {' all': '全国', ' hubei': '湖北', 'other': '湖北外'}

    def three_f(self):
        # 进行三次样条拟合
        model_ls = []
        for square in self.dic.keys():
            Y = self.df[square]
            t = spi.splrep(self.X, Y, k=3)  # 样本点导入，生成参数
            model_ls.append((self.dic[square], t))
            x = np.linspace(1, 30, 1000)
            y = spi.splev(x, t)  # 根据观测点和样条参数，生成插值
            plt.plot(x, y, label=f'三次样条插值曲线-{self.dic[square]}')
        plt.xticks(range(1, 31), self.df['#No. date'])
        plt.legend(loc=2)
        plt.savefig('三次样条插值.jpg')
        plt.show()

        self.predict_f(model_ls, lag=False)

    def lag_f(self):
        # 拉格朗日插值
        model_ls = []
        for square in self.dic.keys():
            Y = np.array(self.df[square])
            t = spi.lagrange(self.X, Y)
            model_ls.append((self.dic[square], t))
            x = np.linspace(1, 30, 90)
            y = t(x)
            plt.plot(x, y, label=f'拉格朗日插值曲线-{self.dic[square]}')
        plt.xticks(range(1, 31), self.df['#No. date'])
        plt.legend(loc=2)
        plt.savefig('拉格朗日插值.jpg')
        plt.show()

        self.predict_f(model_ls)

    def part_lag_f(self):
        # 分段拉格朗日插值
        model_ls = []
        k = 10  # 分10段
        for square in self.dic.keys():
            x_like, y_like = [], []
            Y = np.array(self.df[square])
            for i in range(1, k+1):
                x_part = self.X[3*(i-1):3*i]
                y_part = Y[3*(i-1):3*i]
                t = spi.lagrange(x_part, y_part)
                if i == 1:
                    x = np.linspace(1, 3*i, 100)
                else:
                    x = np.linspace(3*(i-1), 3*i, 100)
                y = t(x)
                x_like.append(x)
                y_like.append(y)
                if i == k:
                    model_ls.append((self.dic[square], t))
            x_lag = sum([list(j) for j in x_like], [])
            y_lag = sum([list(j) for j in y_like], [])
            plt.plot(x_lag, y_lag, label=f'分段拉格朗日插值曲线-{self.dic[square]}')
        plt.xticks(range(1, 31), self.df['#No. date'])
        plt.legend(loc=2)
        plt.savefig('分段拉格朗日插值.jpg')
        plt.show()

        self.predict_f(model_ls)

    def predict_f(self, model_ls, lag=True):
        # 预测
        for model in model_ls:
            x_pre = [31, 32]
            if lag:
                y_pre = model[1](x_pre)
            else:
                y_pre = spi.splev(x_pre, model[1])
            for n, i in enumerate(y_pre):
                print(f'{model[0]}2月{n + 25}日预测值:{int(i)}')


def change_type(args):
    return '0'+str(args)


Insert().three_f()
Insert().lag_f()
Insert().part_lag_f()
